@echo off
REM FLUX 24G显存训练脚本 - Windows版本
REM 适用于24GB VRAM的GPU

echo ========================================
echo FLUX 24G显存基础训练脚本 - Windows版本
echo ========================================

REM 设置环境变量
set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
set CUDA_LAUNCH_BLOCKING=1
set TOKENIZERS_PARALLELISM=false

REM 基础路径定义
set WORKSPACE=/workspace
set AI_DIR=%WORKSPACE%/AI
set PYTHON_PATH=%AI_DIR%/python/bin/python
set SCRIPT_PATH=%AI_DIR%/ComfyUI/custom_nodes/comfyui_lora_train/sd_scripts/flux_train_network.py

REM 模型路径定义
set MODELS_DIR=%AI_DIR%/ComfyUI/models
set FLUX_MODEL=%MODELS_DIR%/checkpoints/flux1-dev-fp8.safetensors
set CLIP_L_MODEL=%MODELS_DIR%/clip/clip_l.safetensors
set T5XXL_MODEL=%MODELS_DIR%/clip/t5xxl_fp8_e4m3fn.safetensors
set AE_MODEL=%MODELS_DIR%/vae/ae.safetensors

REM 数据集和输出路径
set TRAIN_DATA_DIR=%WORKSPACE%/training_data_all/images
set OUTPUT_DIR=%WORKSPACE%/training_data_all/output
set SAMPLE_PROMPTS=%WORKSPACE%/training_data_all/sample_prompts.txt

echo 检查必要文件...

if not exist "%FLUX_MODEL%" (
    echo 错误: FLUX模型文件不存在: %FLUX_MODEL%
    echo 请确保flux1-dev-fp8.safetensors文件在%MODELS_DIR%/checkpoints/目录
    pause
    exit /b 1
)

if not exist "%CLIP_L_MODEL%" (
    echo 错误: CLIP-L模型文件不存在: %CLIP_L_MODEL%
    echo 请确保clip_l.safetensors文件在%MODELS_DIR%/clip/目录
    pause
    exit /b 1
)

if not exist "%T5XXL_MODEL%" (
    echo 错误: T5XXL模型文件不存在: %T5XXL_MODEL%
    echo 请确保t5xxl_fp8_e4m3fn.safetensors文件在%MODELS_DIR%/clip/目录
    pause
    exit /b 1
)

if not exist "%AE_MODEL%" (
    echo 错误: AE模型文件不存在: %AE_MODEL%
    echo 请确保ae.safetensors文件在%MODELS_DIR%/vae/目录
    pause
    exit /b 1
)

if not exist "%TRAIN_DATA_DIR%" (
    echo 错误: 训练数据目录不存在: %TRAIN_DATA_DIR%
    echo 请确保training_data_all/images目录存在并包含训练图像
    pause
    exit /b 1
)

if not exist "%SAMPLE_PROMPTS%" (
    echo 错误: 采样提示词文件不存在: %SAMPLE_PROMPTS%
    echo 请确保sample_prompts.txt文件在%WORKSPACE%/training_data_all/目录
    pause
    exit /b 1
)

REM 创建输出目录
if not exist "%OUTPUT_DIR%" (
    echo 创建输出目录: %OUTPUT_DIR%
    mkdir "%OUTPUT_DIR%"
)

echo ========================================
echo 训练配置信息
echo ========================================
echo 模型: %FLUX_MODEL%
echo 训练数据: %TRAIN_DATA_DIR%
echo 输出目录: %OUTPUT_DIR%
echo 采样提示词: %SAMPLE_PROMPTS%
echo.

echo ========================================
echo 开始FLUX基础训练
echo ========================================

REM 设置环境变量确保使用正确的Python
set PATH=%AI_DIR%/python/bin;%PATH%

REM 基础训练命令
accelerate launch --mixed_precision bf16 --num_cpu_threads_per_process 1 --num_processes 1 --num_machines 1 --dynamo_backend no "%SCRIPT_PATH%" ^
--pretrained_model_name_or_path "%FLUX_MODEL%" ^
--clip_l "%CLIP_L_MODEL%" ^
--t5xxl "%T5XXL_MODEL%" ^
--ae "%AE_MODEL%" ^
--train_data_dir "%TRAIN_DATA_DIR%" ^
--cache_latents_to_disk ^
--save_model_as safetensors ^
--sdpa ^
--persistent_data_loader_workers ^
--max_data_loader_n_workers 2 ^
--seed 42 ^
--gradient_checkpointing ^
--mixed_precision bf16 ^
--save_precision bf16 ^
--network_module networks.lora_flux ^
--network_dim 4 ^
--network_train_unet_only ^
--optimizer_type adamw8bit ^
--learning_rate 1e-4 ^
--cache_text_encoder_outputs ^
--cache_text_encoder_outputs_to_disk ^
--fp8_base ^
--highvram ^
--max_train_epochs 4 ^
--save_every_n_epochs 1 ^
--output_dir "%OUTPUT_DIR%" ^
--output_name flux-lora-24g ^
--timestep_sampling shift ^
--discrete_flow_shift 3.1582 ^
--model_prediction_type raw ^
--guidance_scale 1.0 ^
--sample_every_n_epochs 1 ^
--sample_prompts "%SAMPLE_PROMPTS%" ^
--sample_sampler ddim

if %ERRORLEVEL% neq 0 (
    echo 错误: 训练失败，退出代码: %ERRORLEVEL%
    pause
    exit /b %ERRORLEVEL%
)

echo ========================================
echo FLUX基础训练完成！
echo ========================================

echo 训练完成！模型文件保存在: %OUTPUT_DIR%/flux-lora-24g/
echo 您可以在ComfyUI中加载生成的LoRA模型进行测试
pause
